#include "python.h"
#include <string.h>


#include "../utils/hptime.h"
#include "tracer.h"

#pragma pack(1)
typedef struct _function_enter_block_t {
	// 2 + 18
	uint16_t    codepoint;     // the codepoint will determine whether
	                           // this is a function or a logline
	usec_t      timestamp;
	uint16_t    argcount;      // number of args passed
	uint64_t    argoffset;     // offset into the values listfile
} function_enter_block_t;

/*
 * retval format:
 * 4 lower bits - retval type
 * 60 upper bits - offset/value
 * ==> retval = (offset << 4) | type
 */

#define FUNCTION_RETVAL_UNKNOWN (0)   // retval is ignored
#define FUNCTION_RETVAL_NONE    (1)   // retval is ignored
#define FUNCTION_RETVAL_TRUE    (2)   // retval is ignored
#define FUNCTION_RETVAL_FALSE   (3)   // retval is ignored
#define FUNCTION_RETVAL_INT     (4)   // retval is the (int)value
#define FUNCTION_RETVAL_OBJID   (5)   // retval is object id
#define FUNCTION_RETVAL_OBJOFF  (6)   // retval is offset into the values listfile
#define FUNCTION_RETVAL_EXC     (7)   // retval is offset into the values listfile

typedef struct _function_leave_block_t {
	// 16 bytes
	usec_t      timestamp;
	uint64_t    retval;
} function_leave_block_t;

typedef struct _logline_t {
	// 2 bytes (+ payload)
	uint16_t    codepoint;
	// 'payload' bytes will be written in the leave blob
} logline_t;

#pragma pack()


int tracer_init(tracer_t * self, const char * codepoints_filename,
		const char * values_filename, const char * tracetree_filename,
		int max_depth)
{
	if (hptime_init() != 0) {
		return -1;
	}
	if (htable_init(&self->codepoints_table, 65536) != 0) {
		return -1;
	}
	if (swriter_init(&self->serializer, 20000) != 0) {
		goto cleanup1;
	}
	if (listfile_open(&self->codepoints_list, codepoints_filename) != 0) {
		goto cleanup2;
	}
	if (listfile_open(&self->values_list, values_filename) != 0) {
		goto cleanup3;
	}
	if (treefile_open(&self->trace_tree, tracetree_filename, max_depth) != 0) {
		goto cleanup4;
	}

	return 0;

cleanup4:
	listfile_close(&self->values_list);
cleanup3:
	listfile_close(&self->codepoints_list);
cleanup2:
	swriter_fini(&self->serializer);
cleanup1:
	htable_fini(&self->codepoints_table);
	return -1;
}

int tracer_fini(tracer_t * self)
{
	int ret = 0;
	if (treefile_close(&self->trace_tree) != 0) {
		ret = -1;
	}
	if (listfile_close(&self->codepoints_list) != 0) {
		ret = -1;
	}
	if (htable_fini(&self->codepoints_table) != 0) {
		ret = -1;
	}
	return ret;
}

#define swriter_dump_PyStr(sw, obj) swriter_dump_pstr(sw, PyString_AS_STRING(obj), PyString_GET_SIZE(obj))
#define DUMP_CSTR(sw, obj) if (swriter_dump_cstr(sw, obj) != 0) return -1
#define DUMP_PYSTR(sw, obj) if (swriter_dump_PyStr(sw, obj) != 0) return -1
#define DUMP_UI8(sw, obj) if (swriter_dump_uint8(sw, obj) != 0) return -1
#define DUMP_UI16(sw, obj) if (swriter_dump_uint16(sw, obj) != 0) return -1
#define DUMP_UI32(sw, obj) if (swriter_dump_uint32(sw, obj) != 0) return -1

static inline int _codepoint_dump_codeobject(swriter_t * sw, const PyCodeObject * codeobj)
{
	int i;
	int flags = 0;
	size_t varnames_size = PyTuple_Size(codeobj->co_varnames);
	PyObject * obj;
	flags |= (codeobj->co_flags & CO_VARARGS) ? CODEPOINT_FLAGS_VARARGS : 0;
	flags |= (codeobj->co_flags & CO_VARKEYWORDS) ? CODEPOINT_FLAGS_KWARGS : 0;

	DUMP_UI8(sw, CODEPOINT_TYPE_CODEOBJ);
#ifdef TRACER_DUMP_ABSPATH
	char * realpath = canonicalize_file_name(PyString_AS_STRING(codeobj->co_filename));
	if (realpath == NULL) {
		// fallback to given path
		DUMP_PYSTR(sw, codeobj->co_filename);
	}
	else {
		DUMP_CSTR(sw, realpath);
		free(realpath);
	}
#else
	DUMP_PYSTR(sw, codeobj->co_filename);
#endif
	DUMP_PYSTR(sw, codeobj->co_name);
	DUMP_UI32(sw, codeobj->co_firstlineno);
	DUMP_UI16(sw, codeobj->co_argcount);
	DUMP_UI8(sw, (uint8_t)flags);
	DUMP_UI16(sw, varnames_size);

	for (i = 0; i < varnames_size; i++) {
		obj = PyTuple_GetItem(codeobj->co_varnames, i);
		DUMP_PYSTR(sw, obj);
	}

	return 0;
}

static inline int _codepoint_dump_cfunc(swriter_t * sw, const PyCFunctionObject * cfunc)
{
	DUMP_UI8(sw, CODEPOINT_TYPE_CFUNC);
	if (cfunc->m_module == NULL) {
		DUMP_CSTR(sw, "");
	}
	else {
		PyObject * tmp = PyObject_Str(cfunc->m_module);
		if (tmp == NULL) {
			return -1;
		}
		int ret = swriter_dump_PyStr(sw, tmp);
		Py_DECREF(tmp);
		if (ret != 0) {
			return -1;
		}
	}
	DUMP_CSTR(sw, cfunc->m_ml->ml_name);

	return 0;
}

static inline int _codepoint_dump_strobject(swriter_t * sw, const PyObject * strobj)
{
	DUMP_UI8(sw, CODEPOINT_TYPE_LOG);
	DUMP_PYSTR(sw, strobj);
	return 0;
}

#undef DUMP_UI32
#undef DUMP_UI16
#undef DUMP_UI8
#undef DUMP_PYSTR
#undef DUMP_CSTR
#undef swriter_dump_PyStr

static inline int _codepoint_hash(void * obj)
{
	return (int)(((intptr_t)obj) >> 3);
}

#define GET_CODEPOINT_WITH(sw, obj, dumper)											\
	int index = -1;																	\
	htable_value_t index2;															\
	int hash = _codepoint_hash(obj);													\
																					\
	if (htable_get(&self->codepoints_table, hash, (intptr_t)obj, &index2) == 0) {	\
		return index2;																\
	}																				\
	if (swriter_clear(sw) != 0) {													\
		return -1;																	\
	}																				\
	if (dumper(sw, obj) != 0) {														\
		return -1;																	\
	}																				\
	index = listfile_append(&self->codepoints_list, 									\
			swriter_get_buffer(sw), swriter_get_length(sw));							\
	if (index < 0) {																	\
		return -1;																	\
	}																				\
	if (htable_set(&self->codepoints_table, hash, (intptr_t)obj,						\
			(uint16_t)index) != 0) {													\
		return -1;																	\
	}																				\
	return index;																	\


static inline int _tracer_get_func_codepoint(tracer_t * self, PyCodeObject * codeobj)
{
	GET_CODEPOINT_WITH(&self->serializer, codeobj, _codepoint_dump_codeobject)
}

static inline int _tracer_get_cfunc_codepoint(tracer_t * self, PyCFunctionObject * cfunc)
{
	GET_CODEPOINT_WITH(&self->serializer, cfunc, _codepoint_dump_cfunc)
}

static inline int _tracer_get_log_codepoint(tracer_t * self, PyObject * strobj)
{
	GET_CODEPOINT_WITH(&self->serializer, strobj, _codepoint_dump_strobject)
}

#undef GET_CODEPOINT_WITH

#define PRINT(msg, obj) \
	printf(msg " "); \
	PyObject_Print((PyObject*)obj, stdout, 0); \
	printf("\n")


int tracer_pyfunc_enter(tracer_t * self, PyCodeObject * codeobj, PyObject ** args, int argcount)
{
	function_enter_block_t block;
	int cp = _tracer_get_func_codepoint(self, codeobj);

	if (cp < 0) {
		return -1;
	}

	//PRINT(">>>", codeobj);

	// XXX Py_BEGIN_ALLOW_THREADS ??
	block.codepoint = cp;
	block.argcount = argcount;
	block.argoffset = 0;
	block.timestamp = hptime_get_time();

	return treefile_enter(&self->trace_tree, &block, sizeof(block));
	// XXX Py_END_ALLOW_THREADS ??
}

int tracer_cfunc_enter(tracer_t * self, PyCFunctionObject * cfunc)
{
	function_enter_block_t block;
	int cp = _tracer_get_cfunc_codepoint(self, cfunc);

	if (cp < 0) {
		return -1;
	}

	//PRINT(">>>", cfunc);

	block.codepoint = cp;
	block.argcount = 0;
	block.argoffset = 0;
	block.timestamp = hptime_get_time();

	return treefile_enter(&self->trace_tree, &block, sizeof(block));
}

int tracer_pyfunc_return(tracer_t * self, PyObject * retval)
{
	function_leave_block_t block;

	//printf("<<< pyret\n");

	block.timestamp = hptime_get_time();
	if (retval == NULL) {
		block.retval = FUNCTION_RETVAL_UNKNOWN;
	}
	else if (retval == Py_None) {
		block.retval = FUNCTION_RETVAL_NONE;
	}
	else {
		block.retval = FUNCTION_RETVAL_OBJOFF;
	}

	return treefile_leave(&self->trace_tree, &block, sizeof(block));
}

int tracer_cfunc_return(tracer_t * self)
{
	function_leave_block_t block;

	//printf("<<< cret\n");

	block.timestamp = hptime_get_time();
	// XXX maybe have a cfunc leave block type? to save space?
	// but what about exceptions?
	block.retval = FUNCTION_RETVAL_UNKNOWN;

	return treefile_leave(&self->trace_tree, &block, sizeof(block));
}

int tracer_function_raise(tracer_t * self, PyObject * excinfo)
{
	function_leave_block_t block;

	//printf("<<< cexp\n");

	block.timestamp = hptime_get_time();
	block.retval = FUNCTION_RETVAL_EXC;

	return treefile_leave(&self->trace_tree, &block, sizeof(block));
}

int tracer_log(tracer_t * self, PyObject * fmtstring, PyObject * argstuple)
{
	int ret = -1;
	PyObject * args_as_string;
	int cp = _tracer_get_log_codepoint(self, fmtstring);

	//PRINT("LOG", fmtstring);

	if (cp < 0) {
		return -1;
	}
	uint16_t cp2 = (uint16_t)cp;

	args_as_string = PyObject_Str(argstuple);
	if (args_as_string == NULL) {
		return -1;
	}

	if (treefile_enter(&self->trace_tree, &cp2, sizeof(cp2)) != 0) {
		goto cleanup;
	}
	if (treefile_leave(&self->trace_tree, PyString_AS_STRING(args_as_string),
			(size_t)PyString_GET_SIZE(args_as_string)) != 0) {
		goto cleanup;
	}
	ret = 0;

cleanup:
	Py_DECREF(args_as_string);
	return ret;
}



